import tensorflow as tf
from tensorflow.keras import layers, models
import pandas as pd
from sklearn.preprocessing import MultiLabelBinarizer

# 加载数据集
movies_df = pd.read_csv('movies.csv')

# 处理电影类型数据
genres = movies_df['genres'].apply(lambda x: x.split('|'))
mlb = MultiLabelBinarizer()
genres_encoded = mlb.fit_transform(genres)


from sklearn.model_selection import train_test_split

# 划分训练集和测试集
X_train, X_test, y_train, y_test = train_test_split(genres_encoded, genres_encoded[:, :, 0], test_size=0.2, random_state=42)

# 创建电影的特征输入
movie_input = tf.keras.Input(shape=(genres_encoded.shape[1],), name='Movie-Input')

# linear 是一个线性层，也就是没有激活函数的全连接层（Dense层）。它直接处理输入数据，尝试学习输入与输出之间的线性关系。
# fc1 是一个全连接层，它接受 linear 层的输出，并应用 relu 激活函数。
# fc2 是另一个全连接层，继续处理数据。
# dropout 是一个 Dropout 层，它随机“丢弃”一些神经元的输出，以减少模型过拟合。
# fc3 是最后一个全连接层，进一步处理特征。
# 创建网络层 这里使用顺序网络结构
# 线性层
linear = layers.Dense(units=genres_encoded.shape[1], activation=None)(movie_input)
# 全连接层
fc1 = layers.Dense(units=128, activation='relu')(linear)
fc2 = layers.Dense(units=64, activation='relu')(fc1)
# 增加一个Dropout层来防止过拟合
dropout = layers.Dropout(0.5)(fc2)
fc3 = layers.Dense(units=32, activation='relu')(dropout)

# 输出层
output = layers.Dense(units=genres_encoded.shape[1], activation='sigmoid')(fc3)

# 创建模型
model = models.Model(inputs=movie_input, outputs=output)
model.compile(optimizer='adam', loss='binary_crossentropy', metrics=['accuracy', tf.keras.metrics.AUC(name='auc')])

# 训练模型

tensorboard_callback = tf.keras.callbacks.TensorBoard(log_dir='./logs5', histogram_freq=1)
model.fit(X_train, y_train, validation_data=(X_test, y_test), epochs=10, callbacks=[tensorboard_callback])




import numpy as np
# 推荐函数
def recommend(movie_id, top_n=5):
    # 我们需要使用np.expand_dims来增加输入数据的维度
    movie_vec = genres_encoded[movie_id]
    movie_vec = np.expand_dims(movie_vec, axis=0)  # 这将数据的形状从(18,)转变为(1, 18)
    prediction = model.predict(movie_vec)
    prediction = np.squeeze(prediction)  # 把结果的维度从(1, 18)降低回(18,)

    # 获取排序后的最高评分电影的索引
    sorted_indexes = prediction.argsort()[::-1][:top_n]

    # 返回对应的电影标题
    recommended_movie_ids = movies_df.iloc[sorted_indexes]['movieId']
    return movies_df[movies_df['movieId'].isin(recommended_movie_ids)]['title']


# import pdb
# pdb.set_trace()
# 推荐与movieId为1的电影相似的5部电影
print(recommend(1))
